import os
import sys
import json

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms, datasets
from tqdm import tqdm
import torchvision
from torch.utils.data import DataLoader
from modelv2 import MobileNetV2



def main():

    epochs = 100

    transform = transforms.Compose(
        [transforms.ToTensor(), transforms.Resize((224, 224)), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
    train_set = torchvision.datasets.CIFAR10("../dataset", train=True, transform=transform, download=True)
    val_set = torchvision.datasets.CIFAR10("../dataset", train=False, transform=transform, download=True)

    train_loader = DataLoader(train_set, batch_size=512)
    val_loader = DataLoader(val_set, batch_size=512)
    val_num=len(val_set)

    gpus = [0, 1]  # GPU索引
    device = torch.device("cuda:0")

    # create model
    net = MobileNetV2(num_classes=10)

    # load pretrain weights
    # download url: https://download.pytorch.org/models/mobilenet_v2-b0353104.pth
    # if os.path.exists("./MobileNetV2.pth"):
    #     model_weight_path = "./MobileNetV2.pth"
    # else:
    model_weight_path = "./mobilenet_v2-b0353104.pth"
    assert os.path.exists(model_weight_path), "file {} dose not exist.".format(model_weight_path)
    pre_weights = torch.load(model_weight_path, map_location='cpu')

    # delete classifier weights
    pre_dict = {k: v for k, v in pre_weights.items() if net.state_dict()[k].numel() == v.numel()}
    missing_keys, unexpected_keys = net.load_state_dict(pre_dict, strict=False)



    # freeze features weights 冻结特征提取部分的权重，因为是迁移学习，只需要把在ImageNet上的1000分类转换为我们需要的5分类即可
    for param in net.features.parameters():
        param.requires_grad = False

    net = nn.DataParallel(net, device_ids=gpus)
    net = net.to(gpus[0])  # 主模型

    # define loss function
    loss_function = nn.CrossEntropyLoss()

    # construct an optimizer
    params = [p for p in net.parameters() if p.requires_grad]
    optimizer = optim.Adam(params, lr=0.0001)

    best_acc = 0.0
    save_path = './MobileNetV2_test.pth'
    train_steps = len(train_loader)
    for epoch in range(epochs):
        # train
        net.train()
        running_loss = 0.0
        train_bar = tqdm(train_loader, file=sys.stdout)
        for step, data in enumerate(train_bar):
            images, labels = data
            optimizer.zero_grad()
            logits = net(images.to(device))
            loss = loss_function(logits, labels.to(device))
            loss.backward()
            optimizer.step()

            # print statistics
            running_loss += loss.item()

            train_bar.desc = "train epoch[{}/{}] loss:{:.3f}".format(epoch + 1,
                                                                     epochs,
                                                                     loss)

        # validate
        net.eval()
        acc = 0.0  # accumulate accurate number / epoch
        with torch.no_grad():
            val_bar = tqdm(val_loader, file=sys.stdout)
            for val_data in val_bar:
                val_images, val_labels = val_data
                outputs = net(val_images.to(device))
                # loss = loss_function(outputs, test_labels)
                predict_y = torch.max(outputs, dim=1)[1]
                acc += torch.eq(predict_y, val_labels.to(device)).sum().item()

                val_bar.desc = "valid epoch[{}/{}]".format(epoch + 1,
                                                           epochs)
        val_accurate = acc / val_num
        print('[epoch %d] train_loss: %.3f  val_accuracy: %.3f' %
              (epoch + 1, running_loss / train_steps, val_accurate))

        if val_accurate > best_acc:
            best_acc = val_accurate
            torch.save(net.state_dict(), save_path)

    print('Finished Training')


if __name__ == '__main__':
    main()
